import time
import csv
import torch
from torch.autograd import Variable
from torchvision import models
from DataPreview import dataloader, image_datasets

print('\n==========程序开始运行=============')
train_losses = []  # 存储训练集损失
train_accs = []  # 存储训练集准确率
valid_losses = []  # 存储验证集损失
valid_accs = []  # 存储验证集准确率

print('1. 数据集加载成功，预处理完毕')

model = models.vgg16(pretrained=True)
for param in model.parameters():
    param.requires_grad = False
model.classifier = torch.nn.Sequential(
    torch.nn.Linear(25088, 4096),
    torch.nn.ReLU(),
    torch.nn.Dropout(p=0.5),
    torch.nn.Linear(4096, 4096),
    torch.nn.ReLU(),
    torch.nn.Dropout(p=0.5),
    torch.nn.Linear(4096, 2))
print('2. 模型创建成功')

loss_f = torch.nn.CrossEntropyLoss()
print('3. 目标函数创建完成')

optimizer = torch.optim.Adam(model.classifier.parameters(), lr=0.00001)
print('4. Adam优化器创建完成')

epoch_n = 1
time_open = time.time()
print('5. 超参数设定完毕')

print('\n==========正在测试GPU是否可用==============')
Use_gpu = torch.cuda.is_available()
if Use_gpu:
    print('GPU可用！')
    model = model.cuda()
else:
    print('GPU不可用！')

print('\n==========模型开始训练==============')
for epoch in range(epoch_n):
    print("\n训练批次：Epoch:{}/{}".format(epoch + 1, epoch_n))
    print("-" * 10)

    epoch_start_time = time.time()  # 记录epoch开始时间

    for phase in ["train", "valid"]:
        if phase == "train":
            print("【Train...】")
            model.train(True)
        else:
            print("【Valid...】")
            model.train(False)
        running_loss = 0.0
        running_corrects = 0

        for batch, data in enumerate(dataloader[phase], 1):
            batch_start_time = time.time()  # 记录batch开始时间
            X, y = data
            if Use_gpu:
                X, y = Variable(X.cuda()), Variable(y.cuda())
            else:
                # X, y = Variable(X), Variable(y)
                X, y = X, y
            y_pred = model(X)
            _, pred = torch.max(y_pred.data, 1)
            optimizer.zero_grad()
            loss = loss_f(y_pred, y)
            if phase == "train":
                loss.backward()
                optimizer.step()

            running_loss += loss.data
            running_corrects += torch.sum(pred == y.data)

            if batch % 10 == 0 and phase == "train":
                batch_end_time = time.time()  # 记录batch结束时间
                batch_time = batch_end_time - batch_start_time  # 计算batch耗时
                total_time = time.time() - time_open  # 当前总耗时
                print("Batch {}, Train Loss:{:.4f}, Train ACC:{:.4f}, "
                      "本次10Batch Time:{:.2f}秒，目前已用时{:.2f}秒"
                      .format(batch, running_loss / batch, 100 * running_corrects /
                              (16 * batch), batch_time, total_time))

        epoch_loss = running_loss * 16 / len(image_datasets[phase])
        epoch_acc = 100 * running_corrects / len(image_datasets[phase])
        print("第{}个epoch的{}的训练结果是： Loss:{:.4f} ACC:{:.4f}%".format(epoch+1, phase, epoch_loss, epoch_acc))

        if phase == "train":
            train_losses.append(epoch_loss)
            train_accs.append(epoch_acc)
        else:
            valid_losses.append(epoch_loss)
            valid_accs.append(epoch_acc)
        epoch_end_time = time.time()  # 记录epoch结束时间
        epoch_time = epoch_end_time - epoch_start_time  # 计算epoch耗时
        print("第{}个epoch耗时：{:.2f}秒".format(epoch+1, epoch_time))

time_end = time.time() - time_open
print("总耗时：{:.2f}秒".format(time_end))

# 保存数据到CSV文件
csv_filename = "VGG16_results.csv"

with open(csv_filename, mode='w', newline='') as file:
    writer = csv.writer(file)
    writer.writerow(['Epoch', 'Train Loss', 'Train Accuracy', 'Valid Loss', 'Valid Accuracy'])

    for epoch in range(epoch_n):
        writer.writerow([epoch + 1, train_losses[epoch], train_accs[epoch], valid_losses[epoch], valid_accs[epoch]])

print("数据已保存到文件:", csv_filename)
exit()